from solvers.mpnn import MPNNModel
from tensorflow import keras
from data.bbbp import read_from
from data.smiles_datasets import MPNNDataset, dataset_triplets
from data.smiles_representations import molecule_from_smiles
from rdkit.Chem.Draw import MolsToGridImage
from rdkit import RDLogger
import tensorflow as tf
import warnings
import numpy as np

warnings.filterwarnings("ignore")
RDLogger.DisableLog("rdApp.*")

np.random.seed(42)
tf.random.set_seed(42)

if __name__ == '__main__':
  print('### [INFO] prepare dataset')
  df = read_from('/home/yangw/samples/bbbp.feather')
  print(' ## [INFO] split dataset')
  x_train, y_train, train_index, x_valid, y_valid, valid_index, x_test, y_test, test_index = dataset_triplets(df)
  test_dataset = MPNNDataset(x_test, y_test)

  print('### [INFO] setup model')
  mpnn = MPNNModel(
    atom_dim=29, #x_train[0][0][0].shape[0],
    bond_dim=7, #x_train[1][0][0].shape[0],
  )
  mpnn.compile(
    loss=keras.losses.BinaryCrossentropy(),
    optimizer=keras.optimizers.Adam(learning_rate=5e-4),
    metrics=[keras.metrics.AUC(name="AUC")],
  )
  mpnn.load_weights('weights/mpnn_bbbp_e40.h5')

  # prepare 
  molecules = [molecule_from_smiles(df.smiles.values[index]) for index in test_index]
  y_true = [df.p_np.values[index] for index in test_index]
  y_pred = tf.squeeze(mpnn.predict(test_dataset), axis=1)

  legends = [f"y_true/y_pred = {y_true[i]}/{y_pred[i]:.2f}" for i in range(len(y_true))]
  imgs = MolsToGridImage(molecules, molsPerRow=4, legends=legends)
  imgs.save('test.png')